#include "tools.cpp"

class Solution {
  public:
    vector<int> getLeastNumbers(vector<int>& arr, int k) {
        std::vector<int> result(k);
        top_k_qsort(arr, 0, arr.size() - 1, k);
        for (int i = 0; i < k; ++i) result[i] = arr[i];
        return result;
    }

    void top_k_qsort(std::vector<int>& arr, int start, int end, const int k) {
        int mid = arr[start];
        int low = start + 1, high = end;
        while (high >= low) {
            while (low <= high and arr[low] <= mid) ++low;
            while (low <= high and arr[high] > mid) --high;
            if (low > high) break;
            std::swap(arr[low], arr[high]);
        }
        std::swap(arr[high], arr[start]);
        print(arr);
        if (high == k - 1) return;
        if (high < k - 1) top_k_qsort(arr, high + 1, end, k);
        if (high > k - 1) top_k_qsort(arr, start, high - 1, k);
    }
};

int main() {
    // std::vector<int> arr{3, 12, 4, 5, 2, 1, 3, 42, 3, 2, 123};
    std::vector<int> arr{3, 2, 1};
    Solution s;
    print(s.getLeastNumbers(arr, 2));
    return 0;
}
